from mcp.server.fastmcp import FastMCP
from typing import List
import numpy as np
from scipy import stats

mcp = FastMCP("T检验统计服务")

@mcp.tool()
def one_sample_t_test(sample: List[float], test_value: float) -> str:
    """单样本T检验
    Args:
        sample: 样本数据
        test_value: 需要检验的均值
    Returns:
        检验结果字符串
    """
    try:
        sample_array = np.array(sample)
        k2, p_normal = stats.normaltest(sample_array)
        if p_normal < 0.05:
            norm_warning = "警告: 样本可能不满足正态分布，T检验结果可能不准确。"
        else:
            norm_warning = ""
        t_stat, p_value = stats.ttest_1samp(sample_array, test_value, alternative='two-sided')
        return f"{norm_warning}单样本T检验结果:\nt = {t_stat:.4f}\np = {p_value:.4f} (双尾检验，95%置信水平)"
    except Exception as e:
        return f"计算失败: {str(e)}"

@mcp.tool()
def one_sample_z_test(sample: List[float], population_mean: float, population_std: float) -> str:
    """单样本Z检验
    Args:
        sample: 样本数据
        population_mean: 总体均值（检验值）
        population_std: 总体标准差（已知）
    Returns:
        检验结果字符串
    """
    try:
        arr = np.array(sample)
        n = len(arr)
        if n < 1:
            return "错误: 样本量不足"
        sample_mean = np.mean(arr)
        z = (sample_mean - population_mean) / (population_std / np.sqrt(n))
        # 双尾检验
        p = 2 * (1 - stats.norm.cdf(abs(z)))
        return f"单样本Z检验结果:\nz = {z:.4f}\np = {p:.4f} (双尾检验，95%置信水平)"
    except Exception as e:
        return f"计算失败: {str(e)}"

@mcp.tool()
def paired_sample_t_test(sample1: List[float], sample2: List[float]) -> str:
    """配对样本T检验
    Args:
        sample1: 配对样本1
        sample2: 配对样本2
    Returns:
        检验结果字符串
    """
    try:
        arr1 = np.array(sample1)
        arr2 = np.array(sample2)
        if len(arr1) != len(arr2):
            return "错误: 两个样本长度不一致"
        diff = arr1 - arr2
        k2, p_normal = stats.normaltest(diff)
        if p_normal < 0.05:
            norm_warning = "警告: 差值可能不满足正态分布，配对T检验结果可能不准确。"
        else:
            norm_warning = ""
        t_stat, p_value = stats.ttest_rel(arr1, arr2, alternative='two-sided')
        return f"{norm_warning}配对样本T检验结果:\nt = {t_stat:.4f}\np = {p_value:.4f} (双尾检验)"
    except Exception as e:
        return f"计算失败: {str(e)}"

@mcp.tool()
def two_independent_samples_z_test(sample1: List[float], sample2: List[float], population_std1: float, population_std2: float) -> str:
    """两独立样本Z检验
    Args:
        sample1: 样本1
        sample2: 样本2
        population_std1: 样本1对应总体标准差（已知）
        population_std2: 样本2对应总体标准差（已知）
    Returns:
        检验结果字符串
    """
    try:
        arr1 = np.array(sample1)
        arr2 = np.array(sample2)
        n1 = len(arr1)
        n2 = len(arr2)
        if n1 < 1 or n2 < 1:
            return "错误: 两个样本量均需大于0"
        mean1 = np.mean(arr1)
        mean2 = np.mean(arr2)
        se = np.sqrt(population_std1 ** 2 / n1 + population_std2 ** 2 / n2)
        z = (mean1 - mean2) / se
        # 双尾检验
        p = 2 * (1 - stats.norm.cdf(abs(z)))
        return f"两独立样本Z检验结果:\nz = {z:.4f}\np = {p:.4f} (双尾检验，95%置信水平)"
    except Exception as e:
        return f"计算失败: {str(e)}"

@mcp.tool()
def independent_sample_t_test(sample1: List[float], sample2: List[float]) -> str:
    """独立样本T检验
    Args:
        sample1: 独立样本1
        sample2: 独立样本2
    Returns:
        检验结果字符串
    """
    try:
        arr1 = np.array(sample1)
        arr2 = np.array(sample2)
        k2_1, p_normal_1 = stats.normaltest(arr1)
        k2_2, p_normal_2 = stats.normaltest(arr2)
        if p_normal_1 < 0.05 or p_normal_2 < 0.05:
            norm_warning = "警告: 至少有一个样本可能不满足正态分布，独立样本T检验结果可能不准确。"
        else:
            norm_warning = ""
        stat, p_levene = stats.levene(arr1, arr2)
        equal_var = p_levene >= 0.05
        t_stat, p_value = stats.ttest_ind(arr1, arr2, equal_var=equal_var, alternative='two-sided')
        return f"{norm_warning}独立样本T检验结果:\nt = {t_stat:.4f}\np = {p_value:.4f} (双尾检验)\n方差齐性检验p值 = {p_levene:.4f} (equal_var={equal_var})"
    except Exception as e:
        return f"计算失败: {str(e)}"

if __name__ == "__main__":
    mcp.run(transport="stdio")